package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"

	"unvarnished_transmission/server/config"
)

var (
	conf = config.Config()
)

type ProxyServer struct {
	aListener, bListener *net.TCPListener
	aConns, bConns       chan net.Conn
	wg                   sync.WaitGroup
}

func NewProxyServer() (*ProxyServer, error) {
	aPortA := strings.Split(conf.Server.ProxyAddr, ":")

	aListener, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   net.ParseIP(aPortA[0]),
		Port: mustParsePort(aPortA[1]),
	})

	if err != nil {
		return nil, err
	}

	bPortA := strings.Split(conf.Server.TunnelAddr, ":")

	bListener, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   net.ParseIP(bPortA[0]),
		Port: mustParsePort(bPortA[1]),
	})

	if err != nil {
		return nil, err
	}

	return &ProxyServer{
		aListener: aListener,
		bListener: bListener,
		aConns:    make(chan net.Conn),
		bConns:    make(chan net.Conn),
	}, nil
}

func mustParsePort(port string) int {
	p, err := strconv.Atoi(port)

	if err != nil {
		log.Fatalf("Invalid port number: %s", port)
	}

	return p
}

func (s *ProxyServer) Run() {
	go s.acceptConnections(s.aListener, s.aConns)
	go s.acceptConnections(s.bListener, s.bConns)

	go func() {
		for {
			select {
			case aConn := <-s.aConns:
				bConn := <-s.bConns

				s.wg.Add(2)
				go s.handleConnection(aConn, bConn)
				go s.handleConnection(bConn, aConn)
			case bConn := <-s.bConns:
				aConn := <-s.aConns

				s.wg.Add(2)
				go s.handleConnection(aConn, bConn)
				go s.handleConnection(bConn, aConn)
			}
		}
	}()
}

func (s *ProxyServer) acceptConnections(listener *net.TCPListener, conns chan net.Conn) {
	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			log.Printf("Error accepting connection: %v", err)
			continue
		}

		conns <- conn
	}
}

func (s *ProxyServer) handleConnection(dst, src net.Conn) {
	defer s.wg.Done()
	defer dst.Close()
	defer src.Close()

	io.Copy(dst, src)
}

func main() {
	proxyServer, err := NewProxyServer()
	if err != nil {
		log.Fatalf("Error creating proxy server: %v", err)
	}
	defer proxyServer.aListener.Close()
	defer proxyServer.bListener.Close()

	fmt.Printf("Proxy server started. Listening on ports %s and %s.\n", conf.Server.ProxyAddr, conf.Server.TunnelAddr)

	proxyServer.Run()

	signalChan := make(chan os.Signal, 1)

	<-signalChan
}
